use std::collections::VecDeque;
use dam::utility_contexts::*;
use dam::{simulation::*, templates::pcu::*};
use dam::structures::Time;
use dam::templates::dram::{DRAMConfig, DRAMWriteBundle, DRAMReadBundle, DRAM};
use dam::{
    channel::utils::{EventTime, Peekable},
    channel::{ChannelElement, Receiver, Sender},
    context::Context,
    templates::{
        datastore::Behavior,
        pmu::{PMUReadBundle, PMUWriteBundle, PMU},
        ops::*,
    },
    types::{DAMType, IndexLike},
};

pub struct SambaNova40L {
    // ....
}

pub fn pmu_test() {
    let mut parent = ProgramBuilder::default();
    let size = 1024;
    const CHAN_SIZE: usize = 64;
    const TEST_SIZE: usize = CHAN_SIZE;

    let mut pmu = PMU::<u16, u16, bool>::new(
        TEST_SIZE,
        Behavior {
            mod_address: false,
            use_default_value: false,
        },
    );

    let (wr_addr_send, wr_addr_recv) = parent.bounded(CHAN_SIZE);
    let (wr_data_send, wr_data_recv) = parent.bounded(CHAN_SIZE);
    let (wr_ack_send, wr_ack_recv) = parent.bounded(CHAN_SIZE);

    let write_addr_gen = GeneratorContext::new(
        move || (0..TEST_SIZE).map(|x| u16::try_from(x).unwrap()),
        wr_addr_send,
    );
    parent.add_child(write_addr_gen);

    let wr_data_gen = GeneratorContext::new(
        move || (0..TEST_SIZE).map(|x| u16::try_from(x).unwrap()),
        wr_data_send,
    );

    parent.add_child(wr_data_gen);
    pmu.add_writer(PMUWriteBundle {
        addr: wr_addr_recv,
        data: wr_data_recv,
        ack: wr_ack_send,
    });

    let (rd_addr_send, rd_addr_recv) = parent.bounded(CHAN_SIZE);
    let (rd_data_send, rd_data_recv) = parent.bounded(CHAN_SIZE);
    pmu.add_reader(PMUReadBundle {
        addr: rd_addr_recv,
        resp: rd_data_send,
    });

    let mut rd_addr_gen = FunctionContext::new();
    wr_ack_recv.attach_receiver(&rd_addr_gen);
    rd_addr_send.attach_sender(&rd_addr_gen);
    rd_addr_gen.set_run(move |time| {
        for ind in 0..TEST_SIZE {
            wr_ack_recv.dequeue(time).unwrap();
            let send_time = time.tick();
            rd_addr_send
                .enqueue(
                    time,
                    ChannelElement {
                        time: send_time,
                        data: u16::try_from(ind).unwrap(),
                    },
                )
                .unwrap();
        }
    });

    parent.add_child(rd_addr_gen);
    let checker = CheckerContext::new(
        || (0..TEST_SIZE).map(|x| u16::try_from(x).unwrap()),
        rd_data_recv,
    );
    parent.add_child(checker);
    parent.add_child(pmu);

    let initialized = parent.initialize(InitializationOptions::default()).unwrap();
    let summary = initialized.run(RunOptionsBuilder::default().build().unwrap());
    println!("[PMU] Simulation {}. Elapsed cycles: {:?}", 
                if summary.passed() { "passed" } else { "failed" }, 
                summary.elapsed_cycles()
            );
}

pub fn pcu_test() {
    let mut parent = ProgramBuilder::default();
    let size = 1024;

    const CHAN_SIZE: usize = 8;
    let ingress_op = PCU::READ_ALL_INPUTS;
    let egress_op = PCU::WRITE_ALL_RESULTS;

    let mut pcu = PCU::<u64>::new(
        PCUConfig {
            pipeline_depth: 3,
            num_registers: 3,
        },
        ingress_op,
        egress_op,
    );

    pcu.push_stage(PipelineStage {
        op: ALUMulOp(),
        forward: vec![(2, 1)],
        prev_register_ids: vec![0, 1],
        next_register_ids: vec![],
        output_register_ids: vec![0],
    });

    pcu.push_stage(PipelineStage {
        op: ALUAddOp(),
        forward: vec![],
        prev_register_ids: vec![0, 1],
        next_register_ids: vec![],
        output_register_ids: vec![0],
    });

    let (arg1_send, arg1_recv) = parent.bounded(CHAN_SIZE);
    let (arg2_send, arg2_recv) = parent.bounded(CHAN_SIZE);
    let (arg3_send, arg3_recv) = parent.bounded(CHAN_SIZE);
    let (pcu_out_send, pcu_out_recv) = parent.bounded(CHAN_SIZE);

    pcu.add_input_channel(arg1_recv);
    pcu.add_input_channel(arg2_recv);
    pcu.add_input_channel(arg3_recv);
    pcu.add_output_channel(pcu_out_send);

    let gen1 = GeneratorContext::new(|| (0..size), arg1_send);
    let gen2 = GeneratorContext::new(|| ((size)..(2 * size)), arg2_send);
    let gen3 = GeneratorContext::new(|| ((2 * size)..(3 * size)), arg3_send);
    let checker = CheckerContext::new(
        || (0..size).map(|x| x * (x + size) + (x + size * 2)),
        pcu_out_recv,
    );

    // Run
    parent.add_child(gen1);
    parent.add_child(gen2);
    parent.add_child(gen3);
    parent.add_child(pcu);
    parent.add_child(checker);

    // Type of 'summary': dam::simulation::executed::Executed
    let summary = parent.initialize(InitializationOptions::default())
                        .unwrap()
                        .run(RunOptionsBuilder::default().build().unwrap());

    // End
    println!("[PCU] Simulation {}. Elapsed cycles: {:?}", 
                if summary.passed() { "passed" } else { "failed" }, 
                summary.elapsed_cycles()
            );
}

pub fn dram_test() {
    const TEST_SIZE: usize = 128;
    const NUM_WRITERS: usize = 4;
    const WORK_PER_WRITER: usize = TEST_SIZE / NUM_WRITERS;

    let mut dram = DRAM::<u16, u16, bool>::new(
        DRAMConfig {
            num_simultaneous_requests: 2,
            bandwidth_in_bits: 8,
            latency: Time::new(100),
            capacity: TEST_SIZE,
        },
        Behavior {
            mod_address: false,
            use_default_value: false,
        },
    );

    let mut parent = ProgramBuilder::default();
    let mut ack_channels = Vec::<Receiver<bool>>::with_capacity(NUM_WRITERS);

    (0..NUM_WRITERS).for_each(|split_ind| {
        let low = WORK_PER_WRITER * split_ind;
        let high = low + WORK_PER_WRITER;
        let (addr_send, addr_recv) = parent.bounded(128);
        let (data_send, data_recv) = parent.bounded(128);
        let (ack_send, ack_recv) = parent.bounded(128);
        let (size_send, size_recv) = parent.bounded(128);
        // Address Generator
        let addr_gen =
            GeneratorContext::new(move || [u16::try_from(low).unwrap()].into_iter(), addr_send);
        let data_gen = GeneratorContext::new(
            move || (low..high).map(|x| u16::try_from(x).unwrap()),
            data_send,
        );
        let size_gen = GeneratorContext::new(
            || [u16::try_from(WORK_PER_WRITER).unwrap()].into_iter(),
            size_send,
        );
        dram.add_writer(DRAMWriteBundle {
            addr: addr_recv,
            data: data_recv,
            request_size: size_recv,
            ack: ack_send,
        });
        ack_channels.push(ack_recv);
        parent.add_child(addr_gen);
        parent.add_child(data_gen);
        parent.add_child(size_gen);
    });

    // Create a node that waits for all of the acks to come back, and then issues reads
    let (rd_addr_send, rd_addr_recv) = parent.bounded(128);
    let (rd_data_send, rd_data_recv) = parent.bounded(128);
    let (req_size_send, req_size_recv) = parent.bounded(128);
    let mut read_issue = FunctionContext::new();
    ack_channels.iter_mut().for_each(|chn| {
        chn.attach_receiver(&read_issue);
    });
    rd_addr_send.attach_sender(&read_issue);
    read_issue.set_run(move |time| {
        ack_channels.iter_mut().for_each(|ack| {
            ack.dequeue(time).unwrap();
        });
        let send_time = time.tick();
        rd_addr_send
            .enqueue(
                time,
                ChannelElement {
                    time: send_time,
                    data: 0,
                },
            )
            .unwrap();
    });
    parent.add_child(read_issue);

    let size_issue = GeneratorContext::new(
        || [u16::try_from(TEST_SIZE).unwrap()].into_iter(),
        req_size_send,
    );

    parent.add_child(size_issue);
    dram.add_reader(DRAMReadBundle {
        addr: rd_addr_recv,
        req_size: req_size_recv,
        data: rd_data_send,
    });

    let checker = CheckerContext::new(
        || (0..TEST_SIZE).map(|x| u16::try_from(x).unwrap()),
        rd_data_recv,
    );
    parent.add_child(checker);
    parent.add_child(dram);

    let summary = parent.initialize(InitializationOptions::default())
        .unwrap().run(RunOptions::default());
    
    // End
    println!("[DRAM] Simulation {}. Elapsed cycles: {:?}", 
                if summary.passed() { "passed" } else { "failed" }, 
                summary.elapsed_cycles()
            );
}

fn main() {
    
    // Print Hello
    println!("Welcome to SambaNova 40L with DAM-RS!");

    // Build the system
    pcu_test();
    pmu_test();
    dram_test();
}